# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of Cybertron package.
#
# The Cybertron is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Cybertron
"""

import time
from distutils.version import LooseVersion

from .cybertron import Cybertron
from .model import SchNet, MolCT
from .embedding import MolEmbedding, ConformationEmbedding
from .readout import NodeReadout, AtomwiseReadout, GraphReadout
# from .checkpoint import save_checkpoint, load_checkpoint, load_param_into_net


def _pytorch_version_check():
    """
    Do the PyTorch version check for Cybertron. If PyTorch
    can not be imported, it will raise ImportError. If its
    version is not compatible with current Cybertron version,
    it will print a warning.

    Raise:
        ImportError: If PyTorch can not be imported.
    """
    try:
        import torch
        from torch import __version__ as torch_version
    except ImportError:
        raise ImportError("Can not find PyTorch in current environment. Please install "
                         "PyTorch before using PyTorch Cybertron, by following "
                         "the instruction at https://pytorch.org/get-started/locally/")

    required_torch_version = '1.8.0'
    print(f"Current PyTorch version is {torch_version}")

    if LooseVersion(torch_version) < LooseVersion(required_torch_version):
        print(f"WARNING: Current version of PyTorch is not compatible with Cybertron. "
              f"Some functions might not work or even raise error. Please install PyTorch "
              f"version >= {required_torch_version}. For more details about dependency setting, "
              f"please check the instructions at PyTorch official website https://pytorch.org/")
        warning_countdown = 3
        for i in range(warning_countdown, 0, -1):
            print(f"Please pay attention to the above warning, countdown: {i}")
            time.sleep(1)

_pytorch_version_check()
